from omas.omas_machine import machine_to_omas
from omas.machine_mappings.d3d import thomson_scattering_hardware
from omas import ODS
from matplotlib import pyplot as plt
ods = ODS()
shot6 = 201734
thomson_scattering_hardware(ods, shot6)

eshot = 20173402
options = {'EFIT_tree':  'EFIT'}#, "EFIT_run_id": run_id}
machine_to_omas(ods, 'd3d',eshot, "equilibrium.time", 
                               options=options)
machine_to_omas(ods, 'd3d',eshot, "equilibrium.time_slice.*", 
                               options=options)
machine_to_omas(ods, 'd3d',eshot, "wall.*", options=options)

pshot = 201734002
poptions = {'PROFILES_tree':'OMFIT_PROFS'}
machine_to_omas(ods, 'd3d',pshot, "core_profiles.*", options=poptions)
fig = plt.figure(figsize=[18,5])
ods.plot_equilibrium_summary_and_quality_and_core_profiles(fig=fig)
fig.tight_layout(rect=[0,0,1,0.9])
plt.show()